{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c52e30d1-cb29-4e70-af4a-9c953fcb0f2e",
   "metadata": {},
   "source": [
    "# Quickstart: Vector search using Azure OpenAI Embeddings and Elasticsearch\n",
    "\n",
    "This tutorial demonstrates how to use the [Azure OpenAI API](https://azure.microsoft.com/en-in/products/ai-services/openai-service) to create [embeddings](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/embeddings?tabs=console) and store them in Elasticsearch. Elasticsearch will enable us to perform vector search (Knn) to find similar documents."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88303061-f357-43d8-8b63-c4f79e9a1746",
   "metadata": {},
   "source": [
    "## setup\n",
    "\n",
    "* Elastic Credentials - Create [Cloud deployment](https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud) to get all Elastic credentials (`ELASTIC_CLOUD_ID`, `ELASTIC_API_KEY`).\n",
    "\n",
    "* `AZURE_OPENAI_API_KEY` - To use the Azure OpenAI API, you need an API key. [Follow](https://learn.microsoft.com/en-us/azure/ai-services/openai/quickstart?tabs=command-line%2Cpython-new&pivots=programming-language-python#retrieve-key-and-endpoint) to create a key.\n",
    "* `AZURE_OPENAI_ENDPOINT` - Endpoint for your Azure OpenAI Resource.\n",
    "* `AZURE_DEPLOYMENT_ID` - The deployment name you chose when you deployed the model.\n",
    "* `AZURE_OPENAI_API_VERSION` - The API version to use for this operation. This follows the YYYY-MM-DD format."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76ca723c-6148-4682-a5ae-486e73cb2b94",
   "metadata": {},
   "source": [
    "## Install packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef1f1e52-f892-489f-8947-3e4698f5f5c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install -q -U openai elasticsearch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d86d3fa-4ca0-41b6-a4bc-81bacf26bf02",
   "metadata": {},
   "source": [
    "## Import packages and credentials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb62d8fb-6c34-44fd-bc94-18b644422ee8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import AzureOpenAI\n",
    "from elasticsearch import Elasticsearch, helpers\n",
    "from getpass import getpass\n",
    "import os\n",
    "\n",
    "ELASTIC_API_KEY = getpass(\"Elastic API Key :\")\n",
    "ELASTIC_CLOUD_ID = getpass(\"Elastic Cloud ID :\")\n",
    "\n",
    "AZURE_OPENAI_API_KEY = getpass(\"Azure OpenAI API Key :\")\n",
    "AZURE_OPENAI_ENDPOINT = getpass(\"Azure OpenAI Endpoint :\")\n",
    "AZURE_DEPLOYMENT_ID = getpass(\"Azure Deployment ID :\")\n",
    "AZURE_OPENAI_API_VERSION = getpass(\"Azure OpenAI API Version :\")\n",
    "\n",
    "elastic_index_name = \"azure-openai-vector-search-demo\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b22dc16-c0a0-48f0-979d-5d21c17bd264",
   "metadata": {},
   "source": [
    "## Embedding generation\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca56532d-7c82-4e2b-aecf-2173520d3696",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_embeddings(text):\n",
    "    client = AzureOpenAI(\n",
    "        api_key=AZURE_OPENAI_API_KEY,\n",
    "        api_version=AZURE_OPENAI_API_VERSION,\n",
    "        azure_endpoint=AZURE_OPENAI_ENDPOINT,\n",
    "    )\n",
    "\n",
    "    response = client.embeddings.create(\n",
    "        input=text,\n",
    "        model=AZURE_DEPLOYMENT_ID,\n",
    "    )\n",
    "\n",
    "    return response.data[0].embedding\n",
    "\n",
    "\n",
    "sample_text = \"India generally experiences a hot summer from March to June, with temperatures often exceeding 40°C in central and northern regions. Monsoon season, from June to September, brings heavy rainfall, especially in the western coast and northeastern areas. Post-monsoon months, October and November, mark a transition with decreasing rainfall. Winter, from December to February, varies in temperature across the country, with colder conditions in the north and milder weather in the south. India's diverse climate is influenced by its geographical features, resulting in regional\"\n",
    "embeddings = generate_embeddings(sample_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6239eda7-3bed-43dd-a6a8-a8369b907d5c",
   "metadata": {},
   "source": [
    "## Connecting Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cbade18-3049-46f1-8d3e-5b22d4aade5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "es = Elasticsearch(cloud_id=ELASTIC_CLOUD_ID, api_key=ELASTIC_API_KEY)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20d070c8-9e19-48a3-bc3b-5f22067eb63f",
   "metadata": {},
   "source": [
    "## Index document with Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e02ca81e-7caa-4505-95c6-3c6be7843c8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc = {\"text\": sample_text, \"text_embedding\": embeddings}\n",
    "\n",
    "resp = es.index(index=elastic_index_name, document=doc)\n",
    "\n",
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afa0d371-afbf-4f98-9cd1-ee457839f323",
   "metadata": {},
   "source": [
    "## Searching for document with Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d71eeacc-d0c8-4035-b052-a1c03300aec0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "ID: SxtQyY4BMvvuJ06pSACG\n",
      "\n",
      "Text: India generally experiences a hot summer from March to June, with temperatures often exceeding 40°C in central and northern regions. Monsoon season, from June to September, brings heavy rainfall, especially in the western coast and northeastern areas. Post-monsoon months, October and November, mark a transition with decreasing rainfall. Winter, from December to February, varies in temperature across the country, with colder conditions in the north and milder weather in the south. India's diverse climate is influenced by its geographical features, resulting in regional\n"
     ]
    }
   ],
   "source": [
    "q = \"How's weather in India?\"\n",
    "\n",
    "embeddings = generate_embeddings(q)\n",
    "\n",
    "resp = es.search(\n",
    "    index=elastic_index_name,\n",
    "    knn={\n",
    "        \"field\": \"text_embedding\",\n",
    "        \"query_vector\": embeddings,\n",
    "        \"k\": 10,\n",
    "        \"num_candidates\": 100,\n",
    "    },\n",
    ")\n",
    "\n",
    "\n",
    "for result in resp[\"hits\"][\"hits\"]:\n",
    "    pretty_output = f\"\\n\\nID: {result['_id']}\\n\\nText: {result['_source']['text']}\"\n",
    "    print(pretty_output)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
